#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import os
import sys
import getopt
import subprocess

usage = f"""使用方式：
1、vim此脚本并配置class Cfg内的参数配置。如安装位置、备机数量、端口号、ip等。

2、支持的操作：

    - 安装主备机群：  python3 gsom.py init
        根据配置，在一台机器上安装出一个主备集群。每个节点的端口号都是配置的端口号递增一段距离。
    
    - 查看安装信息：  python3 gsom.py show
    
    - 启动主备集群：  python3 gsom.py start
    
    - 关闭主备集群：  python3 gsom.py stop
    
    - 删除主备集群：  python3 gsom.py delete

    - 连接某个节点:   python3 gsom.py c|connect 节点编号如123456等

3、快速命令
    alias gsom='python3 {sys.argv[0]}'

"""


class Cfg(object):
    gausshome = os.environ['GAUSSHOME']                 # 环境变量GAUSSHOME
    ip = '127.0.0.1'                                    # 本机ip
    port = 20088                                        # db运行端口
    count = 2                                           # 备机数量
    dnpath = os.path.join(gausshome, 'gsom')       # 安装位置

    simple_log = True                                   # 简化日志打印，不打印例如gs_initdb等命令的输出，只打印命令内容。

    @staticmethod
    def init():
        if Cfg.gausshome is None or Cfg.ip is None or Cfg.port is None or Cfg.count is None or Cfg.dnpath is None:
            print('先配置class Cfg的参数。')
            exit(1)

        if not os.access(Cfg.dnpath, os.F_OK):
            Cmd.execute('mkdir ' + Cfg.dnpath)


class Dn(object):

    @staticmethod
    def construct_dninfos():
        dn_list = []
        for i in range(0, Cfg.count + 1):
            dn = Dn()
            dn.nodename = 'dn{}'.format(i + 1)
            dn.init_mode = 'primary' if i == 0 else 'standby'
            dn.path = os.path.join(Cfg.dnpath, dn.nodename)
            dn.port = Cfg.port + i * 100
            dn.localport = dn.port + 10
            dn.localheartbeatport = dn.port + 20
            
            dn_list.append(dn)
        
        # replconninfo
        for i, dn in enumerate(dn_list):
            replconninfos = []
            
            for j, rdn in enumerate(dn_list):
                if i == j:
                    continue
                replconninfos.append(
                    'localhost={0} localport={1} localheartbeatport={2} remotehost={3} remoteport={4} remoteheartbeatport={5}'.format(
                        Cfg.ip, dn.localport, dn.localheartbeatport,
                        Cfg.ip, rdn.localport, rdn.localheartbeatport
                    )
                )
            
            dn.replconninfos = replconninfos
        
        return dn_list
        
    def __init__(self):
        self.nodename = None
        self.init_mode = None
        self.path = None
        self.port = None
        
        self.replconninfos = []
        self.localport = None
        self.localheartbeatport = None
    
    def __str__(self):
        return 'nodename:{0}  init-mode:{1}  path:{2}  port:{3}'.format(
            self.nodename,
            self.init_mode,
            self.path,
            self.port
        )
    
    def getpid(self):
        pmpid = os.path.join(self.path, 'postmaster.pid')
        if not os.access(pmpid, os.F_OK):
            return
        with open(pmpid, 'r') as f:
            line = f.readline()
            return int(line.strip())  # todo, 未运行则返回None
    
    def setguc(self, name, value):
        cmd = """gs_guc set -D {0} -c "{1}='{2}'" """.format(self.path, name, value)
        Cmd.execute(cmd)
        
    def start(self, mode=None, just_get_cmd=False):
        mode = self.init_mode if mode is None else mode
        cmd = "gs_ctl start -D {0} -M {1}".format(self.path, mode)
        if just_get_cmd:
            return cmd
        Cmd.execute(cmd)

    def initdb(self, just_get_cmd=False):
        cmd = "gs_initdb -D {0} -w FastInit@123 --nodename='{1}'".format(self.path, self.nodename)
        if just_get_cmd:
            return cmd
        Cmd.execute(cmd)

    def stop(self, just_get_cmd=False):
        cmd = "gs_ctl stop -D {0}".format(self.path)
        if just_get_cmd:
            return cmd
        Cmd.execute(cmd)

    def restart(self, just_get_cmd=False):
        cmd = "gs_ctl restart -D {0}".format(self.path)
        if just_get_cmd:
            return cmd
        Cmd.execute(cmd)

    def build(self, just_get_cmd=False):
        cmd = "gs_ctl build -D {0}".format(self.path)
        if just_get_cmd:
            return cmd
        Cmd.execute(cmd)

    def delete(self):
        pid = self.getpid()
        if pid is not None:
            Cmd.execute(f'kill -9 {pid}', ignore_err=True)
        Cmd.execute(f'rm -fr {self.path}')

    def connect(self):
        cmd = f'gsql -p {self.port} -d postgres -r'
        print(cmd)
        os.system(cmd)


class Cmd(object):

    @staticmethod
    def execute(cmd, exec_async=False, ignore_err=False):
        def _exec_sync():
            print("exec cmd:", cmd)
            output = None
            if Cfg.simple_log:
                state, output = subprocess.getstatusoutput(cmd)
            else:
                state = os.system(cmd)

            if state != 0:
                if output is not None:
                    print(output)
                print(f"执行失败，退出码为 {state}")
                if not ignore_err:
                    exit(state)

        def _exec_async():
            print("exec cmd async:", cmd)
            if Cfg.simple_log:
                return subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
            else:
                return subprocess.Popen(cmd, shell=True)

        if exec_async:
            return _exec_async()
        else:
            _exec_sync()
            return None
    
    @staticmethod
    def execute_group(cmds, exec_parallel=False, ignore_err=False):
        processes = []
        for cmd in cmds:
            p = Cmd.execute(cmd, exec_parallel, ignore_err)
            processes.append(p)

        if not exec_parallel:
            return

        err_ocurr = False
        for i, process in enumerate(processes):
            stdout, stderr = process.communicate()
            exit_code = process.returncode

            if exit_code != 0:
                err_ocurr = True
                print(f"有失败的子进程{i}，退出码为 {exit_code}。子进程输出如下：")
                print('标准输出:', stdout.decode('utf-8'))
                print('错误输出:', stderr.decode('utf-8'))

        if not ignore_err and err_ocurr:
            exit(1)


def init(dninfos):
    # initdb
    cmds = []
    for dn in dninfos:
        cmds.append(dn.initdb(just_get_cmd=True))
    Cmd.execute_group(cmds, exec_parallel=True)
    
    # set guc
    for dn in dninfos:
        dn.setguc('port', dn.port)
        dn.setguc('max_wal_senders', 100)
        dn.setguc('application_name', dn.nodename)
        for i, replconninfo in enumerate(dn.replconninfos):
            dn.setguc(f'replconninfo{i + 1}', replconninfo)
    
    # start up
    cmds = []
    for dn in dninfos:
        cmds.append(dn.start(just_get_cmd=True))
    Cmd.execute_group(cmds, exec_parallel=True)
    
    # build
    cmds = []
    for dn in dninfos:
        if dn.init_mode == 'primary':
            continue
        cmds.append(dn.build(just_get_cmd=True))
    Cmd.execute_group(cmds, exec_parallel=True)
    
    show(dninfos)


def start(dninfos):
    cmds = []
    for dn in dninfos:
        cmds.append(dn.start(just_get_cmd=True))
    Cmd.execute_group(cmds)


def restart(dninfos):
    cmds = []
    for dn in dninfos:
        cmds.append(dn.restart(just_get_cmd=True))
    Cmd.execute_group(cmds)


def stop(dninfos):
    cmds = []
    for dn in dninfos:
        cmds.append(dn.stop(just_get_cmd=True))
    Cmd.execute_group(cmds)


def delete(dninfos):
    for dn in dninfos:
        dn.delete()
    print('删除完成。')


def show(dninfos):
    print('安装信息：')
    for dn in dninfos:
        print('    ', dn)
    print('')
    print('快速命令：')
    print(f"    alias gsom='python3 {sys.argv[0]}'")


def connect(dninfos, i=1):
    assert 0 < i <= len(dninfos), '没这号节点'
    dn = dninfos[i - 1]
    print(f'连接:', dn)
    dn.connect()


if __name__ == '__main__':
    Cfg.init()
    dninfos = Dn.construct_dninfos()
    
    if len(sys.argv) == 1:
        print(usage)
        exit(0)
    elif sys.argv[1] == 'init':
        init(dninfos)
    elif sys.argv[1] == 'show':
        show(dninfos)
    elif sys.argv[1] == 'start':
        start(dninfos)
    elif sys.argv[1] == 'restart':
        restart(dninfos)
    elif sys.argv[1] == 'stop':
        stop(dninfos)
    elif sys.argv[1] == 'delete':
        delete(dninfos)
    elif sys.argv[1] in ['connect', 'c']:
        connect(dninfos, int(sys.argv[2]))
    else:
        print(usage)
        exit(0)
